# charged_particle_simulator_final_formatted.py
import pygame
import math
import importlib

# --- 1. Initialization and Setup ---
try:
    import config
except ImportError:
    print("错误: 未找到 'config.py' 文件。请确保它与此脚本位于同一目录中。")
    exit()

pygame.init()
SCREEN_WIDTH, SCREEN_HEIGHT = 1200, 800
UI_PANEL_WIDTH = 280

# --- Colors and Fonts ---
BG_COLOR = (15, 20, 30)
UI_BG_COLOR = (30, 35, 45)
BORDER_COLOR = (80, 80, 90)
TEXT_COLOR = (230, 230, 230)
HOVER_COLOR = (70, 90, 120)
COLOR_POSITIVE = (0, 255, 255)
COLOR_NEGATIVE = (255, 0, 128)
COLOR_VELOCITY = (50, 255, 50)
COLOR_FORCE = (255, 80, 80)
FIELD_E_COLOR = (*COLOR_POSITIVE, 20)
FIELD_M_COLOR = (*COLOR_NEGATIVE, 20)
FIELD_SYMBOL_COLOR = (150, 150, 150, 100)

screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Physics Simulator (Final Version)")
font = pygame.font.Font(None, 28)
small_font = pygame.font.Font(None, 22)
large_font = pygame.font.Font(None, 80)
vector_font = pygame.font.SysFont("dejavusans", 20, bold=True)


# --- 2. Helper Functions ---
def draw_text(text, font, color, surface, x, y, center=False):
    textobj = font.render(text, True, color)
    textrect = textobj.get_rect()
    if center:
        textrect.center = (x, y)
    else:
        textrect.topleft = (x, y)
    surface.blit(textobj, textrect)


def draw_arrow(surface, color, start, end, width=2, arrow_size=12, label=None):
    """
    绘制一个更形象、更具视觉冲击力的箭头。
    - 箭头翼更宽，头部更大。
    - 标签位置经过微调，以更好地适应新的箭头形状。
    """
    try:
        start_v = pygame.math.Vector2(start)
        end_v = pygame.math.Vector2(end)
        
        # 绘制平滑的箭头主干
        pygame.draw.aaline(surface, color, start_v, end_v, width)

        # 如果箭头太短，就不画头部，避免视觉混乱
        if (start_v - end_v).length() < arrow_size:
            return

        # 计算箭头方向和角度
        # 我们从终点指向起点计算角度，这样更容易处理
        angle = math.atan2(start_v.y - end_v.y, start_v.x - end_v.x)
        
        # 定义箭头翼的角度，pi/7 大约是 25.7 度，比之前的 30 度稍窄，但配合更大的arrow_size效果更好
        wing_angle = math.pi / 7
        
        # 计算箭头翼的两个端点
        p1 = (
            end_v.x + arrow_size * math.cos(angle - wing_angle),
            end_v.y + arrow_size * math.sin(angle - wing_angle),
        )
        p2 = (
            end_v.x + arrow_size * math.cos(angle + wing_angle),
            end_v.y + arrow_size * math.sin(angle + wing_angle),
        )
        
        # 绘制实心、平滑的三角形箭头头部
        pygame.draw.polygon(surface, color, (end_v, p1, p2))

        # 如果有标签，绘制它
        if label:
            # 将标签放置在箭头头部后方更远一点的位置，避免重叠
            label_offset = arrow_size * 1.5
            # 角度需要反转，因为我们是从终点计算的
            label_angle = angle + math.pi
            
            label_pos_x = end_v.x + label_offset * math.cos(label_angle)
            label_pos_y = end_v.y + label_offset * math.sin(label_angle)
            
            text_surf = vector_font.render(label, True, color)
            text_rect = text_surf.get_rect(center=(label_pos_x, label_pos_y))
            surface.blit(text_surf, text_rect)

    except (ValueError, ZeroDivisionError):
        # 避免在起点和终点重合时崩溃
        pass


# --- 3. Physics & UI Classes ---
class Particle:
    def __init__(self, mass, charge, pos_m, vel_ms, label="P"):
        self.mass, self.charge = mass, charge
        self.pos_m, self.vel_ms = pygame.math.Vector2(pos_m), pygame.math.Vector2(vel_ms)
        self.label = label
        self.pos_px = pygame.math.Vector2(0, 0)
        self.force_N = pygame.math.Vector2(0, 0)
        self.trajectory_px = []
        self.color = COLOR_POSITIVE if charge >= 0 else COLOR_NEGATIVE

    def update(self, dt, simulator):
        def get_derivatives(pos, vel):
            force = simulator.calculate_total_force_at(self, pos, vel)
            acceleration = force / self.mass if self.mass != 0 else pygame.math.Vector2(0, 0)
            return vel, acceleration
        
        k1_vel, k1_acc = get_derivatives(self.pos_m, self.vel_ms)
        k2_vel, k2_acc = get_derivatives(self.pos_m + 0.5 * k1_vel * dt, self.vel_ms + 0.5 * k1_acc * dt)
        k3_vel, k3_acc = get_derivatives(self.pos_m + 0.5 * k2_vel * dt, self.vel_ms + 0.5 * k2_acc * dt)
        k4_vel, k4_acc = get_derivatives(self.pos_m + k3_vel * dt, self.vel_ms + k3_acc * dt)
        
        self.pos_m += (dt / 6.0) * (k1_vel + 2 * k2_vel + 2 * k3_vel + k4_vel)
        self.vel_ms += (dt / 6.0) * (k1_acc + 2 * k2_acc + 2 * k3_acc + k4_acc)

    def update_screen_pos(self, m_per_px, origin_offset_m):
        pos_m_shifted = self.pos_m - origin_offset_m
        self.pos_px.x = pos_m_shifted.x / m_per_px
        self.pos_px.y = SCREEN_HEIGHT - (pos_m_shifted.y / m_per_px)

        # Correctly formatted trajectory logic
        if len(self.trajectory_px) == 0 or self.pos_px.distance_to(self.trajectory_px[-1]) > 1:
            self.trajectory_px.append(self.pos_px.copy())
        
        if len(self.trajectory_px) > 4000:
            self.trajectory_px.pop(0)

    def draw(self, surface):
        pygame.draw.circle(surface, (*self.color, 50), self.pos_px, 8)
        pygame.draw.circle(surface, self.color, self.pos_px, 5)

    def draw_vectors(self, surface, show_vel, show_force):
        vector_display_length = 50
        if show_vel and self.vel_ms.length() > 0:
            vel_px = self.vel_ms.copy()
            vel_px.y *= -1
            end_point = self.pos_px + vel_px.normalize() * vector_display_length
            draw_arrow(surface, COLOR_VELOCITY, self.pos_px, end_point, label='v')
        if show_force and self.force_N.length() > 0:
            force_px = self.force_N.copy()
            force_px.y *= -1
            end_point = self.pos_px + force_px.normalize() * vector_display_length
            draw_arrow(surface, COLOR_FORCE, self.pos_px, end_point, label='F')

    def draw_trajectory(self, surface):
        if len(self.trajectory_px) > 1:
            pygame.draw.aalines(surface, self.color, False, self.trajectory_px)

class Field:
    def __init__(self, shape, region_m):
        self.shape, self.region_m = shape, region_m

    def get_bounding_box(self):
        if self.shape == 'rect':
            x, y, w, h = self.region_m
            return x, y, x + w, y + h
        if self.shape == 'circle':
            cx, cy, r = self.region_m
            return cx - r, cy - r, cx + r, cy + r
        return 0, 0, 0, 0

    def contains_m(self, pos_m):
        if self.shape == 'rect':
            x, y, w, h = self.region_m
            return x <= pos_m.x <= x + w and y <= pos_m.y <= y + h
        if self.shape == 'circle':
            cx, cy, r = self.region_m
            return math.hypot(pos_m.x - cx, pos_m.y - cy) <= r
        return False

    def _convert_region_to_px(self, mpp, offset):
        region_shifted = self.region_m[:]
        if self.shape in ['rect', 'circle']:
            region_shifted[0] -= offset.x
            region_shifted[1] -= offset.y
        
        if self.shape == 'rect':
            x_m, y_m, w_m, h_m = region_shifted
            x_px = x_m / mpp
            w_px = w_m / mpp
            h_px = h_m / mpp
            y_px = SCREEN_HEIGHT - ((y_m + h_m) / mpp)
            return pygame.Rect(x_px, y_px, w_px, h_px)
        
        if self.shape == 'circle':
            cx_m, cy_m, r_m = region_shifted
            cx_px = cx_m / mpp
            cy_px = SCREEN_HEIGHT - (cy_m / mpp)
            r_px = r_m / mpp
            return (int(cx_px), int(cy_px)), int(r_px)
        return None

class ElectricField(Field):
    def __init__(self, shape, region_m, strength_vec):
        super().__init__(shape, region_m)
        self.strength_Vm = pygame.math.Vector2(strength_vec)

    def get_force_on(self, particle, **kwargs):
        if self.contains_m(particle.pos_m):
            return particle.charge * self.strength_Vm
        return pygame.math.Vector2(0, 0)

    def draw(self, surface, mpp, offset):
        region_px = self._convert_region_to_px(mpp, offset)
        if not region_px: return
        
        s = pygame.Surface(surface.get_size(), pygame.SRCALPHA)
        if self.shape == 'rect':
            pygame.draw.rect(s, FIELD_E_COLOR, region_px)
        elif self.shape == 'circle':
            pygame.draw.circle(s, FIELD_E_COLOR, region_px[0], region_px[1])
        surface.blit(s, (0, 0))
        self._draw_field_vectors(surface, region_px)

    def _draw_field_vectors(self, surface, region_px):
        if self.strength_Vm.length() == 0: return

        draw_direction = self.strength_Vm.normalize()
        draw_direction.y *= -1
        
        stepx, stepy, arrow_len = 40, 100, 100

        if self.shape == 'rect':
            for x in range(int(region_px.left), int(region_px.right), stepx):
                for y in range(int(region_px.top), int(region_px.bottom), stepy):
                    start = pygame.math.Vector2(x + stepx//2, y + stepy//2)
                    end = start + draw_direction * arrow_len
                    draw_arrow(surface, FIELD_SYMBOL_COLOR, start, end, 1, 5)
        elif self.shape == 'circle':
            center, radius = region_px
            for x in range(center[0]-radius, center[0]+radius, stepx):
                for y in range(center[1]-radius, center[1]+radius, stepy):
                    start = pygame.math.Vector2(x, y)
                    if math.hypot(x-center[0], y-center[1]) < radius:
                        end = start + draw_direction * arrow_len
                        draw_arrow(surface, FIELD_SYMBOL_COLOR, start, end, 1, 5)

class MagneticField(Field):
    def __init__(self, shape, region_m, strength_T):
        super().__init__(shape, region_m)
        self.strength_T = strength_T

    def get_force_on(self, particle, vel_override=None, **kwargs):
        vel = vel_override if vel_override is not None else particle.vel_ms
        if self.contains_m(particle.pos_m):
            vx, vy = vel.x, vel.y
            return pygame.math.Vector2(particle.charge*vy*self.strength_T, particle.charge*-vx*self.strength_T)
        return pygame.math.Vector2(0, 0)

    def draw(self, surface, mpp, offset):
        region_px = self._convert_region_to_px(mpp, offset)
        if not region_px: return

        s = pygame.Surface(surface.get_size(), pygame.SRCALPHA)
        if self.shape == 'rect': pygame.draw.rect(s, FIELD_M_COLOR, region_px)
        elif self.shape == 'circle': pygame.draw.circle(s, FIELD_M_COLOR, region_px[0], region_px[1])
        surface.blit(s, (0, 0))
        self._draw_field_symbols(surface, region_px)

    def _draw_field_symbols(self, surface, region_px):
        step = 40
        if self.strength_T == 0: return
        draw_func = self._draw_b_dot if self.strength_T > 0 else self._draw_b_cross

        if self.shape == 'rect':
            for x in range(int(region_px.left), int(region_px.right), step):
                for y in range(int(region_px.top), int(region_px.bottom), step):
                    draw_func(surface, (x + step//2, y + step//2))
        elif self.shape == 'circle':
            center, radius = region_px
            for x in range(center[0]-radius, center[0]+radius, step):
                for y in range(center[1]-radius, center[1]+radius, step):
                    if math.hypot(x-center[0], y-center[1]) < radius:
                        draw_func(surface, (x, y))

    def _draw_b_dot(self, s, p):
        pygame.draw.circle(s, FIELD_SYMBOL_COLOR, p, 6, 1)
        pygame.draw.circle(s, FIELD_SYMBOL_COLOR, p, 1)

    def _draw_b_cross(self, s, p):
        pygame.draw.aaline(s, FIELD_SYMBOL_COLOR, (p[0]-4, p[1]-4), (p[0]+4, p[1]+4))
        pygame.draw.aaline(s, FIELD_SYMBOL_COLOR, (p[0]-4, p[1]+4), (p[0]+4, p[1]-4))

class Button:
    def __init__(self, rect, text):
        self.rect, self.text, self.is_hovered = pygame.Rect(rect), text, False
    def handle_event(self, event):
        if event.type == pygame.MOUSEMOTION: self.is_hovered = self.rect.collidepoint(event.pos)
        return event.type == pygame.MOUSEBUTTONDOWN and event.button == 1 and self.is_hovered
    def draw(self, surface):
        color = HOVER_COLOR if self.is_hovered else UI_BG_COLOR
        pygame.draw.rect(surface, color, self.rect, border_radius=5)
        pygame.draw.rect(surface, BORDER_COLOR, self.rect, 1, border_radius=5)
        draw_text(self.text, font, TEXT_COLOR, surface, self.rect.centerx, self.rect.centery, True)

class Checkbox:
    def __init__(self, rect, text, checked=False):
        self.rect, self.text, self.checked, self.is_hovered = pygame.Rect(rect), text, checked, False
    def handle_event(self, event):
        if event.type == pygame.MOUSEMOTION: self.is_hovered = self.rect.collidepoint(event.pos)
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1 and self.is_hovered:
            self.checked = not self.checked
            return True
        return False
    def draw(self, surface):
        color = HOVER_COLOR if self.is_hovered else BORDER_COLOR
        pygame.draw.rect(surface, UI_BG_COLOR, self.rect)
        pygame.draw.rect(surface, color, self.rect, 1, border_radius=3)
        if self.checked:
            p1 = (self.rect.left+4, self.rect.centery)
            p2 = (self.rect.centerx-2, self.rect.bottom-4)
            p3 = (self.rect.right-4, self.rect.top+4)
            pygame.draw.aaline(surface, COLOR_VELOCITY, p1, p2)
            pygame.draw.aaline(surface, COLOR_VELOCITY, p2, p3)
        draw_text(self.text, small_font, TEXT_COLOR, surface, self.rect.right+10, self.rect.centery-small_font.get_height()//2)

class Simulator:
    def __init__(self):
        self.clock = pygame.time.Clock()
        self.is_running, self.was_running = False, False
        self.particles, self.fields = [], []
        self.gravity_ms2 = None
        self.m_per_px, self.time_scaling = 1.0, 1.0
        self.origin_offset_m = pygame.math.Vector2(0, 0)
        self.sim_area_width = SCREEN_WIDTH - UI_PANEL_WIDTH
        self.setup_ui()
        self.reset()

    def setup_ui(self):
        px = self.sim_area_width + 30
        bw = (UI_PANEL_WIDTH - 80) // 2
        y = 40
        self.start_button = Button((px, y, bw, 40), "Start")
        self.pause_button = Button((px + bw + 20, y, bw, 40), "Pause")
        y += 60
        self.reset_button = Button((px, y, bw * 2 + 20, 40), "Reset Scene")
        y1, y2 = 220, 380
        y = y1
        self.cb_trajectory = Checkbox((px, y, 20, 20), "Show Trajectory", True)
        y += 40
        self.cb_velocity = Checkbox((px, y, 20, 20), "Show Velocity", False)
        y += 40
        self.cb_force = Checkbox((px, y, 20, 20), "Show Force", False)
        y = y2
        self.cb_gravity = Checkbox((px, y, 20, 20), "Enable Gravity", False)
        self.ui_elements = [self.start_button, self.pause_button, self.reset_button, self.cb_gravity, self.cb_trajectory, self.cb_velocity, self.cb_force]

    def reset(self):
        self.is_running, self.was_running = False, False
        self.load_config_and_reset()

    def load_config_and_reset(self):
        try:
            importlib.reload(config)
            self.cfg_scale, self.cfg_scene = config.SCALING_CONFIG, config.SCENE_CONFIG
        except Exception as e:
            print(f"Config Error: {e}")
            self.cfg_scale, self.cfg_scene = {}, {}
        self.time_scaling = self.cfg_scale.get("time_scaling", 1.0)
        self.setup_scene()
        self.auto_zoom()

    def setup_scene(self):
        self.particles.clear()
        self.fields.clear()
        scene = self.cfg_scene
        g_cfg = scene.get("gravity", {})
        self.cb_gravity.checked = g_cfg.get("enabled", False)
        self.gravity_ms2 = pygame.math.Vector2(g_cfg.get("g_vector", [0, 0])) if self.cb_gravity.checked else None
        for p_data in scene.get("particles", []):
            self.particles.append(Particle(p_data.get("mass", 1), p_data.get("charge", 0), p_data.get("position", [0, 0]), p_data.get("velocity", [0, 0]), p_data.get("label", "N/A")))
        for f_data in scene.get("fields", []):
            if f_data.get("type") == "electric":
                self.fields.append(ElectricField(f_data.get("shape"), f_data.get("region"), f_data.get("strength_vector", [0, 0])))
            elif f_data.get("type") == "magnetic":
                self.fields.append(MagneticField(f_data.get("shape"), f_data.get("region"), f_data.get("strength", 0)))

    def auto_zoom(self):
        if not self.fields:
            self.m_per_px, self.origin_offset_m = 1.0, pygame.math.Vector2(0, 0)
            return
        min_x, min_y, max_x, max_y = float('inf'), float('inf'), float('-inf'), float('-inf')
        for f in self.fields:
            fxm, fym, fxa, fya = f.get_bounding_box()
            min_x, min_y, max_x, max_y = min(min_x, fxm), min(min_y, fym), max(max_x, fxa), max(max_y, fya)
        
        w_m, h_m = max_x - min_x, max_y - min_y
        if w_m == 0 or h_m == 0:
            # Handle case with point/line fields
            w_m = w_m if w_m > 0 else 1.0
            h_m = h_m if h_m > 0 else 1.0

        pad = self.cfg_scale.get("auto_zoom_padding_factor", 1.2)
        w_m *= pad
        h_m *= pad

        scale_x = w_m / self.sim_area_width if self.sim_area_width > 0 else float('inf')
        scale_y = h_m / SCREEN_HEIGHT if SCREEN_HEIGHT > 0 else float('inf')
        self.m_per_px = max(scale_x, scale_y)
        if self.m_per_px <= 0: self.m_per_px = 1.0 # Safety check

        cx_m = min_x + (max_x - min_x) / 2
        cy_m = min_y + (max_y - min_y) / 2
        sim_cx_m = (self.sim_area_width / 2) * self.m_per_px
        sim_cy_m = (SCREEN_HEIGHT / 2) * self.m_per_px
        self.origin_offset_m = pygame.math.Vector2(cx_m - sim_cx_m, cy_m - sim_cy_m)

    def calculate_total_force_at(self, particle, pos_m, vel_ms):
        total_force = pygame.math.Vector2(0, 0)
        if self.gravity_ms2:
            total_force += self.gravity_ms2 * particle.mass
        dummy = Particle(particle.mass, particle.charge, pos_m, vel_ms)
        for f in self.fields:
            total_force += f.get_force_on(dummy, vel_override=vel_ms)
        return total_force

    def update_logic(self, dt):
        if not self.is_running:
            self.was_running = False
            return
        self.was_running = True
        sim_dt = dt * self.time_scaling
        for p in self.particles:
            p.update(sim_dt, self)
            p.force_N = self.calculate_total_force_at(p, p.pos_m, p.vel_ms)

    def update_screen_positions(self):
        for p in self.particles:
            p.update_screen_pos(self.m_per_px, self.origin_offset_m)

    def run(self):
        running = True
        while running:
            dt_s = self.clock.tick(60) / 1000.0
            if not self.handle_events():
                running = False
            self.update_logic(dt_s)
            self.update_screen_positions()
            self.draw()
        pygame.quit()

    def draw(self):
        screen.fill(BG_COLOR)
        for f in self.fields:
            f.draw(screen, self.m_per_px, self.origin_offset_m)
        for p in self.particles:
            if self.cb_trajectory.checked:
                p.draw_trajectory(screen)
            p.draw(screen)
            p.draw_vectors(screen, self.cb_velocity.checked, self.cb_force.checked)

        pygame.draw.rect(screen, UI_BG_COLOR, (self.sim_area_width, 0, UI_PANEL_WIDTH, SCREEN_HEIGHT))
        pygame.draw.line(screen, BORDER_COLOR, (self.sim_area_width, 0), (self.sim_area_width, SCREEN_HEIGHT))
        self.draw_ui_panel_content()
        if not self.is_running and self.was_running:
            s = pygame.Surface((self.sim_area_width, SCREEN_HEIGHT), pygame.SRCALPHA)
            s.fill((0, 0, 0, 128))
            screen.blit(s, (0, 0))
            draw_text("PAUSED", large_font, (*TEXT_COLOR, 200), screen, self.sim_area_width // 2, SCREEN_HEIGHT // 2, True)
        pygame.display.flip()

    def draw_ui_panel_content(self):
        px = self.sim_area_width + 30
        for el in self.ui_elements:
            el.draw(screen)
        draw_text("Display", font, BORDER_COLOR, screen, px, 180)
        draw_text("Physics", font, BORDER_COLOR, screen, px, 340)
        y = 460
        draw_text(f"Scale: 1px = {self.m_per_px:.2e} m", small_font, BORDER_COLOR, screen, px, y)
        y += 25
        draw_text(f"Time Scale: {self.time_scaling:.2e}", small_font, BORDER_COLOR, screen, px, y)
        y += 40
        if self.particles:
            p = self.particles[0]
            draw_text(f"Info: {p.label}", font, BORDER_COLOR, screen, px, y)
            y += 40
            draw_text(f"Pos:({p.pos_m.x:.2e}, {p.pos_m.y:.2e}) m", small_font, TEXT_COLOR, screen, px, y)
            y += 25
            draw_text(f"Vel:{p.vel_ms.length():.2e} m/s", small_font, TEXT_COLOR, screen, px, y)
            y += 25
            draw_text(f"Force:{p.force_N.length():.2e} N", small_font, TEXT_COLOR, screen, px, y)

    def handle_events(self):
        for e in pygame.event.get():
            if e.type == pygame.QUIT:
                return False
            if self.start_button.handle_event(e):
                self.is_running = True
            if self.pause_button.handle_event(e):
                self.is_running = False
            if self.reset_button.handle_event(e):
                self.reset()
            for cb in self.ui_elements:
                if isinstance(cb, Checkbox) and cb.handle_event(e):
                    if cb == self.cb_gravity:
                        g_cfg = self.cfg_scene.get("gravity", {})
                        g_vec = g_cfg.get("g_vector", [0, 0])
                        self.gravity_ms2 = pygame.math.Vector2(g_vec) if self.cb_gravity.checked else None
        return True

if __name__ == '__main__':
    simulator = Simulator()
    simulator.run()